package main

import (
	"errors"
	"fmt"
	"sync"
	"time"
)

// Snowflake ID 生成器
// 结构：0 - 41位时间戳 - 10位机器ID - 12位序列号
type Snowflake struct {
	mu           sync.Mutex
	epoch        int64 // 起始时间戳（毫秒）
	workerID     int64 // 机器ID (0-31)
	datacenterID int64 // 数据中心ID (0-31)
	sequence     int64 // 序列号 (0-4095)
	lastTime     int64 // 上次生成ID的时间戳
}

const (
	workerIDBits     = 5
	datacenterIDBits = 5
	sequenceBits     = 12

	maxWorkerID     = -1 ^ (-1 << workerIDBits)     // 31
	maxDatacenterID = -1 ^ (-1 << datacenterIDBits) // 31
	maxSequence     = -1 ^ (-1 << sequenceBits)     // 4095

	timeShift       = workerIDBits + datacenterIDBits + sequenceBits // 22
	datacenterShift = workerIDBits + sequenceBits                    // 17
	workerShift     = sequenceBits                                   // 12
)

// NewSnowflake 创建Snowflake ID生成器
func NewSnowflake(datacenterID, workerID int64) (*Snowflake, error) {
	if datacenterID < 0 || datacenterID > maxDatacenterID {
		return nil, errors.New("datacenterID out of range")
	}
	if workerID < 0 || workerID > maxWorkerID {
		return nil, errors.New("workerID out of range")
	}

	return &Snowflake{
		epoch:        1609459200000, // 2021-01-01 00:00:00
		datacenterID: datacenterID,
		workerID:     workerID,
		sequence:     0,
		lastTime:     0,
	}, nil
}

// NextID 生成下一个ID
func (s *Snowflake) NextID() (int64, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	now := time.Now().UnixMilli()

	// 时钟回拨检测
	if now < s.lastTime {
		offset := s.lastTime - now
		if offset <= 5 { // 5ms内等待
			time.Sleep(time.Duration(offset) * time.Millisecond)
			now = time.Now().UnixMilli()
		} else {
			return 0, fmt.Errorf("clock moved backwards by %dms", offset)
		}
	}

	if now == s.lastTime {
		// 同一毫秒内，序列号递增
		s.sequence = (s.sequence + 1) & maxSequence
		if s.sequence == 0 {
			// 序列号溢出，等待下一毫秒
			now = s.waitNextMillis(now)
		}
	} else {
		// 不同毫秒，序列号重置
		s.sequence = 0
	}

	s.lastTime = now

	// 组装ID
	id := ((now - s.epoch) << timeShift) |
		(s.datacenterID << datacenterShift) |
		(s.workerID << workerShift) |
		s.sequence

	return id, nil
}

func (s *Snowflake) waitNextMillis(currentTime int64) int64 {
	for currentTime == s.lastTime {
		time.Sleep(100 * time.Microsecond)
		currentTime = time.Now().UnixMilli()
	}
	return currentTime
}

// Parse 解析 Snowflake ID
func (s *Snowflake) Parse(id int64) map[string]int64 {
	timestamp := (id >> timeShift) + s.epoch
	datacenterID := (id >> datacenterShift) & maxDatacenterID
	workerID := (id >> workerShift) & maxWorkerID
	sequence := id & maxSequence

	return map[string]int64{
		"timestamp":    timestamp,
		"datacenterID": datacenterID,
		"workerID":     workerID,
		"sequence":     sequence,
	}
}

func main() {
	// 创建 Snowflake 生成器（数据中心ID=1，机器ID=1）
	sf, err := NewSnowflake(1, 1)
	if err != nil {
		panic(err)
	}

	fmt.Println("Snowflake ID Generator Demo")
	fmt.Println("=" * 50)

	// 生成10个 ID
	for i := 0; i < 10; i++ {
		id, err := sf.NextID()
		if err != nil {
			fmt.Printf("Error: %v\n", err)
			continue
		}

		fmt.Printf("ID: %d\n", id)

		// 解析 ID
		parsed := sf.Parse(id)
		ts := time.UnixMilli(parsed["timestamp"])
		fmt.Printf("  Time: %s, DC: %d, Worker: %d, Seq: %d\n",
			ts.Format("2006-01-02 15:04:05.000"),
			parsed["datacenterID"],
			parsed["workerID"],
			parsed["sequence"])
	}

	// 性能测试
	fmt.Println("\n性能测试...")
	count := 100000
	start := time.Now()
	for i := 0; i < count; i++ {
		_, err := sf.NextID()
		if err != nil {
			fmt.Printf("Error: %v\n", err)
			break
		}
	}
	elapsed := time.Since(start)
	fmt.Printf("生成 %d 个ID耗时: %v\n", count, elapsed)
	fmt.Printf("QPS: %.0f\n", float64(count)/elapsed.Seconds())
}
